import pandas as pd
import numpy as np
from ortools.linear_solver import pywraplp

#数据加载
data = pd.read_csv('./family_data.csv', index_col = 'family_id')
#print(data)

# Step2，数据预处理
# 1）计算Perference Cost矩阵 pcost_mat
# 2）计算Accounting Cost矩阵 acost_mat
# 3）计算每个家庭的人数 FAMILY_SIZE
# 4）每个家庭的倾向选择（choice_） DESIRED

#n代表家庭成员个数，如果满足第choice需求，需要的penalty
def get_penalty(n, choice):
    if choice == 0:
        penalty  =0
    if choice == 1:
        penalty = 50
    if choice == 2:
        penalty = 50 + 9*n
    if choice == 3:
        penalty = 100 + 9 * n
    if choice == 4:
        penalty = 200 + 9 * n
    if choice == 5:
        penalty = 200 + 18 * n
    if choice == 6:
        penalty = 300 + 18 * n
    if choice == 7:
        penalty = 300 + 36 * n
    if choice == 8:
        penalty = 400 + 36 * n
    if choice == 9:
        penalty = 500 + (36 +199)* n
    if choice >9:
        penalty = 500 + (36 +398)* n

    return penalty

N_DAYS = 100
N_FAMILY = 5000
MIN_OCCUPANCY = 125  #最小承载量
MAX_OCCUPANCY = 300  #最大承载量

#计算pcost_matt,每个家庭，在什么时候(day 0-99)访问时的penalty
#大小5000*100的矩阵
pcost_mat = np.full(shape = (N_FAMILY,100),fill_value=999999)
for f in range(N_FAMILY):
    #家庭成员数
    f_num = data.loc[f,'n_people']
    #print(f_num)

    #对于第f个家庭，初始化pcost=other_choice
    pcost_mat[f,:] = get_penalty(f_num,10)

    #计算choice0-9的penalty
    for choice in range(10):
        temp = data.loc[f][choice]   #choice的天数
        penalty = get_penalty(f_num, choice)
        pcost_mat[f, temp-1] = penalty

#print(pcost_mat)

#计算accouting penalty矩阵，前一天的参观人数，当天的参观人数
acost_mat = np.zeros(shape=(MAX_OCCUPANCY+1,MAX_OCCUPANCY+1),dtype=np.float64)

for i in range(acost_mat.shape[0]):     #当天安排的人数
    for j in range(acost_mat.shape[1]): #前一天安排的人数
        diff = abs(i-j)
        acost_mat[i,j] =max(0, (i-125) / 400 * i **(0.5 + diff/50.0))

#print(acost_mat)
FAMILY_SIZE = data['n_people'].values
#print(FAMILY_SIZE)

#DESIRED代表每个家庭CHOICE-1
#pcost_mat的时间是1-100，对应的小标是0-99
DESIRED=data.values[:,:-1]-1
#print(DESIRED)

# Step3，使用LP和MIP求解 规划方案
# 1）先使用LP 对绝大部分家庭进行规划
# 2）再使用MIP 对剩余家庭进行规划
# 3）汇总两边的结果 => 最终规划方案

#线性规划
def solveLP():
    #线性规划优化器
    solver = pywraplp.Solver('AssignmentProblem', pywraplp.Solver.GLOP_LINEAR_PROGRAMMING)
    x={}     #family_id在第j天是否参观
    #每一天有哪些家庭,有100天，每一天的家庭
    candidates=[[] for x in range(N_DAYS)]

    for i in range(N_FAMILY): #family_id
        for j in DESIRED[i,:]:   #family_id的choice
            candidates[j].append(i)     #在第j天，有第i个family参观
            #定义决策变量x[i,j],i代表family_id, j代表第j天参观
            x[i,j] = solver.BoolVar('x[%i,%i]' % (i,j))
    #print(x)
    #每天参观的人数

    daily_occupancy = [solver.Sum([ x[i, j] * FAMILY_SIZE[i] for i in candidates[j]])  for j in range(N_DAYS)]  #j代表1-100天
    #daily_occupancy = [solver.Sum(x[i,j] * FAMILY_SIZE[i] for i in candidates[j]]) for j in range(N_DAYS)]

    #每个家庭在10个choice中出现的情况
    family_presence= [solver.Sum(x[i,j] for j in DESIRED[i,:]) for i in range(N_FAMILY)]

    #定义目标函数preference cost部分
    preference_cost = solver.Sum(pcost_mat[i,j] * x[i,j] for i in range(N_FAMILY) for j in DESIRED[i,:])

    #满足preference cost最小
    solver.Minimize((preference_cost))

    #人为增加约束条件Constraints
    for j in range(N_DAYS-1):   #j代表当天，j+1代表前一天
        #当前人数不超过前一天人数+25
        solver.Add(daily_occupancy[j] - daily_occupancy[j+1]<=25)
        solver.Add(daily_occupancy[j+1] - daily_occupancy[j] <= 25)
    #每个家庭都在10个choice里出现1次
    for i in range(N_FAMILY):
        solver.Add(family_presence[i] == 1)

    #每天访问人数约束
    for j in range(N_DAYS):
        solver.Add(daily_occupancy[j] >= MIN_OCCUPANCY)
        solver.Add(daily_occupancy[j] <= MAX_OCCUPANCY)

    result = solver.Solve()

    temp = [(i,j,x[i, j].solution_value()) for i in range(N_FAMILY) for j in DESIRED[i, :] if x[i, j].solution_value() > 0]
    #print(temp)
    #得到参观日期的安排
    df = pd.DataFrame(temp, columns=['family_id','day','result'])
    return df, result

df,result = solveLP()

#所有的assign结果
print(df)

THRS = 0.999
#安排上的
assigned_df=df[df['result']>THRS]

#没有被安排上的
unassigned_df=df[(df['result']<THRS) & (df['result']>1-THRS)]

#计算每天访问的人数(根据assigned_df)
assigned_df['family_size'] = FAMILY_SIZE[assigned_df.family_id]

#按照day的方式进行聚合
occupancy=assigned_df.groupby('day').family_size.sum()

print(occupancy)

min_occupancy= np.array([max(0,MIN_OCCUPANCY-x) for x in occupancy])
max_occupancy= np.array([max(0,MAX_OCCUPANCY-x) for x in occupancy])

#使用整数规划进行求解
def solveIP(families, min_occupancy, max_occupancy):
    #创建求解器
    solver = pywraplp.Solver("AssignmentProblem", pywraplp.Solver.CBC_MIXED_INTEGER_PROGRAMMING)
    n_familes = len(families)

    x = {}  # family_id在第j天是否参观
    # 每一天有哪些家庭,有100天，每一天的家庭
    candidates = [[] for x in range(N_DAYS)]

    for i in families:  # family_id
        for j in DESIRED[i, :]:  # family_id的choice
            candidates[j].append(i)  # 在第j天，有第i个family参观
            # 定义决策变量x[i,j],i代表family_id, j代表第j天参观
            x[i, j] = solver.BoolVar('x[%i,%i]' % (i, j))
    # print(x)
    # 每天参观的人数

    daily_occupancy = [solver.Sum([x[i, j] * FAMILY_SIZE[i] for i in candidates[j]]) for j in
                       range(N_DAYS)]  # j代表1-100天
    # daily_occupancy = [solver.Sum(x[i,j] * FAMILY_SIZE[i] for i in candidates[j]]) for j in range(N_DAYS)]

    # 每个家庭在10个choice中出现的情况
    family_presence = [solver.Sum(x[i, j] for j in DESIRED[i, :]) for i in families]

    # 定义目标函数preference cost部分
    preference_cost = solver.Sum(pcost_mat[i, j] * x[i, j] for i in families for j in DESIRED[i, :])

    # 满足preference cost最小
    solver.Minimize((preference_cost))


    # 每个家庭都在10个choice里出现1次
    for i in range(n_familes):
        solver.Add(family_presence[i] == 1)

    # 每天访问人数约束
    for j in range(N_DAYS):
        solver.Add(daily_occupancy[j] >= min_occupancy[j])
        solver.Add(daily_occupancy[j] <= max_occupancy[j])

    result = solver.Solve()

    temp = [(i, j) for i in families for j in DESIRED[i, :] if
            x[i, j].solution_value() > 0]
    # print(temp)

    #计算剩余家庭安排
    df = pd.DataFrame(temp, columns= ['family_id', 'day'])

    return df

unassigned = unassigned_df.family_id.unique()
result = solveIP(unassigned, min_occupancy, max_occupancy)

#print(result)

#数据合并
df = pd.concat((assigned_df[['family_id','day']], result)).sort_values('family_id')
print(df)

#根据安排情况，计算这个安排的preference cost
def pcost(predictions):
    daily_occupancy = np.zeros(N_DAYS+1, dtype=np.int64)
    penalty=0
    for (i,p) in enumerate(predictions):
        #计算家庭人数
        n = FAMILY_SIZE[i]
        #第i个家庭，p天访问时的cost
        penalty +=pcost_mat[i,p]
        #计算当天的人数
        daily_occupancy[p] += n

    return penalty,daily_occupancy

#根据安排情况，计算安排的accounting cost
def acost(daily_occupancy):
    accounting_cost = 0
    num_out_of_range = 0
    for day in range(N_DAYS):
        n_p1 = daily_occupancy[day+1]  #前一天
        n = daily_occupancy[day]       #当天

        #如果超过了承载范围，则设置 out_of_range
        num_out_of_range += (n>MAX_OCCUPANCY) or (n<MIN_OCCUPANCY)

        #计算accounting cost
        accounting_cost += acost_mat[n,n_p1]

    return accounting_cost, num_out_of_range

#根据安排prediction，计算
def cost_function(prediction):
    #基于predcition计算preference cost和accounting cost
    penalty, daily_occupancy = pcost(prediction)  #计算preference cost和每天承载数量
    accounting_cost, num_out_of_range = acost(daily_occupancy) #根据每天承载数量计算accounting cost

    final_score = penalty + accounting_cost + num_out_of_range*99999999   #最后一项是对越界的惩罚
    return final_score

prediction = df.day.values
print(cost_function(prediction))

def save_result(pred, filename):
    result = pd.DataFrame(range(N_FAMILY), columns=['family_id'])
    result['assigned_day'] = pred+1
    result.to_csv(filename, index=False)
    print(filename + ' saved')
    return result

result = save_result(prediction, 'submission1.csv')
